import numpy as np
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.neighbors import KNeighborsRegressor


class MissingValueHandler:
    def __init__(self, data):
        """
        初始化类，传入包含缺失值的DataFrame对象
        """
        self.data = data
        self.missing_values = data.isnull().sum()

    def dropna(self, how='any', axis=0):
        """
        删除含有缺失值的行或列
        """
        return self.data.dropna(how=how, axis=axis)

    def fill_with_constant(self, value, method="constant"):
        """
        使用固定值填充缺失值
        """
        if method == "mean":
            fill_value = self.data.mean()
        elif method == "median":
            fill_value = self.data.median()
        elif method == "mode":
            fill_value = self.data.mode().iloc[0]
        else:
            fill_value = value

        return self.data.fillna(fill_value)

    def fill_with_imputer(self, strategy="mean"):
        """
        使用scikit-learn的SimpleImputer进行插补
        """
        imputer = SimpleImputer(strategy=strategy)
        filled_data = pd.DataFrame(imputer.fit_transform(self.data), columns=self.data.columns)
        return filled_data

    def fill_with_knn(self, n_neighbors=5):
        """
        使用K近邻插补
        """
        # 假设所有数值列都适合进行KNN插补
        numeric_features = self.data.select_dtypes(include=[np.number]).columns.tolist()

        imputer = KNeighborsRegressor(n_neighbors=n_neighbors)
        for feature in numeric_features:
            imp = imputer.fit(self.data.dropna(subset=[feature]), self.data[feature].dropna())
            self.data.loc[self.data[feature].isnull(), feature] = imp.predict(
                self.data.loc[self.data[feature].isnull(), numeric_features])

        return self.data

    def interpolate_linear(self, columns=None):
        """
        对指定列进行线性插值
        """
        if columns is None:
            columns = self.data.columns

        for col in columns:
            self.data[col].interpolate(inplace=True, method='linear')

        return self.data
